
import torch.nn as nn
import torch

netG = nn.Sequential(
        nn.Conv2d(16, 64, kernel_size=3, stride=1, padding=1),
        nn.PixelShuffle(2),
        nn.LeakyReLU(0.1, True)
    )

if isinstance(netG, nn.DataParallel):
    netG = netG.module
netG.load_state_dict(torch.load("./pixel_shuffle.pth"), strict=False)
netG.eval().cuda()
dummy_input = torch.randn(1, 16, 200, 200).cuda()
dynamic_axes = { 
            'color_input':  {2: 'dy_num',3:'dy_num'},
            'color_output':{2: 'dy_num',3:'dy_num'}
            }
torch.onnx.export(netG, dummy_input, r"./pixel_shuffle.onnx",opset_version=10, export_params=True,input_names = ['color_input'],output_names=['color_output'],dynamic_axes=dynamic_axes)
print("      pth convert  onnx completes successfully !!!! ")